// Copyright (c) 2019 Presto Labs Pte. Ltd.
// Author: donggu
#pragma once
#include <algorithm>
#include <ctime>  // strptime
#include <experimental/filesystem>
#include <memory>
#include <regex>
#include <set>
#include <string>
#include <utility>
#include <vector>

#include "coin/feed/fastfeed/feed.pb.h"
#include "orio/io/base.h"
#include "orio/io/concat.h"
#include "orio/record/simple.h"

namespace impl {
using fastfeed::proto::Feed;

class FastFeedProtoReader : public orio::io::Reader<Feed> {
 public:
  FastFeedProtoReader(
      std::unique_ptr<orio::io::Reader<orio::record::SimpleRecord>> reader,
      size_t max_record_size)
      : reader_(std::move(reader)), buf_(new char[max_record_size]) {}
  bool Read(Feed* out) {
    if (!reader_->Read(reinterpret_cast<orio::record::SimpleRecord*>(buf_.get()))) {
      return false;
    }
    auto* r = reinterpret_cast<orio::record::SimpleRecord*>(buf_.get());
    if (!out->ParseFromArray(r->data, r->len)) {
      throw std::runtime_error("[FastFeedProtoReader] fails to parse proto");
    }

    // only the first message have the symbol
    if (out->symbol().empty()) {
      CHECK(!symbol_.empty()) << "No symbol set: " << out->DebugString();
      out->set_symbol(symbol_);
    } else if (symbol_.empty()) {
      symbol_ = out->symbol();
    }

    return true;
  }

 private:
  std::unique_ptr<orio::io::Reader<orio::record::SimpleRecord>> reader_;
  std::unique_ptr<char[]> buf_;
  std::string symbol_;
};

struct MergeEntry {
  size_t idx;
  int64_t timestamp;
  bool friend operator<(const MergeEntry& p, const MergeEntry& q) {
    return std::tie(p.timestamp, p.idx) < std::tie(q.timestamp, q.idx);
  }
};

class MergedFastFeedReader : public orio::io::Reader<Feed> {
 public:
  explicit MergedFastFeedReader(std::vector<std::unique_ptr<orio::io::Reader<Feed>>> readers)
      : readers_(std::move(readers)), buffer_(readers_.size()), last_reader_idx_(-1) {
    // initial read
    for (size_t i = 0; i < readers_.size(); i++) {
      if (readers_[i]->Read(&buffer_[i])) {
        set_.insert(MergeEntry{i, buffer_[i].timestamp()});
      }
    }
  }

  bool Read(Feed* out) override {
    if (set_.empty()) {
      return false;
    } else {
      const MergeEntry& entry = *set_.begin();
      size_t idx = entry.idx;
      // copy record
      out->CopyFrom(buffer_[idx]);
      last_reader_idx_ = idx;
      // pop from set
      set_.erase(set_.begin());
      // read next
      if (readers_[idx]->Read(&buffer_[idx])) {
        set_.insert(MergeEntry{idx, buffer_[idx].timestamp()});
      }
      return true;
    }
  }

 private:
  std::vector<std::unique_ptr<orio::io::Reader<Feed>>> readers_;
  std::vector<Feed> buffer_;
  std::set<MergeEntry> set_;
  int last_reader_idx_;
};

class IntervalFastFeedReader : public orio::io::Reader<Feed> {
 public:
  IntervalFastFeedReader(
      std::unique_ptr<orio::io::Reader<Feed>> reader,
      int64_t begin_timestamp, /* inclusive */
      int64_t end_timestamp)   /* exclusive */
      : reader_(std::move(reader)),
        begin_timestamp_(begin_timestamp),
        end_timestamp_(end_timestamp) {}

  bool Read(Feed* out) override {
    while (reader_->Read(out) && out->timestamp() < end_timestamp_) {
      if (out->timestamp() >= begin_timestamp_ ||
          out->feed_type() == ::fastfeed::proto::FEED_TYPE_METADATA) {
        return true;
      }
    }
    return false;  // EOF
  }

 private:
  std::unique_ptr<orio::io::Reader<Feed>> reader_;
  int64_t begin_timestamp_;
  int64_t end_timestamp_;
};

class UniverseFastFeedReader : public orio::io::Reader<Feed> {
 public:
  explicit UniverseFastFeedReader(std::vector<std::unique_ptr<orio::io::Reader<Feed>>> readers)
      : readers_(std::move(readers)), buffer_(readers_.size()), last_reader_idx_(-1) {
    // initial read
    for (size_t i = 0; i < readers_.size(); i++) {
      if (readers_[i]->Read(&buffer_[i])) {
        set_.insert(MergeEntry{i, buffer_[i].timestamp()});
      }
    }
  }

  bool Read(Feed* out) override {
    if (set_.empty()) {
      return false;
    } else {
      const MergeEntry& entry = *set_.begin();
      size_t idx = entry.idx;
      // copy record
      out->CopyFrom(buffer_[idx]);
      last_reader_idx_ = idx;
      // pop from set
      set_.erase(set_.begin());
      // read next
      if (readers_[idx]->Read(&buffer_[idx])) {
        set_.insert(MergeEntry{idx, buffer_[idx].timestamp()});
      }
      return true;
    }
  }

 private:
  std::vector<std::unique_ptr<orio::io::Reader<Feed>>> readers_;
  std::vector<Feed> buffer_;
  std::set<MergeEntry> set_;
  int last_reader_idx_;
};


}  // namespace impl

using impl::FastFeedProtoReader;
using impl::IntervalFastFeedReader;
using impl::MergedFastFeedReader;
